"""
@File    : captcha_auto_recognizer.py
@Author  : timo.tan
@Time    : 2025/07/28
"""

import cv2
import numpy as np
import matplotlib.pyplot as plt
import os

def preprocess_captcha_image(image_path, debug=True):
    """
    预处理验证码图片以提高识别率
    """
    # 读取图片
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"无法读取图片: {image_path}")
    
    # 保存原始图像尺寸用于调试
    original_shape = img.shape
    print(f"原始图像尺寸: {original_shape}")
    
    # 转换为灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 尝试不同的二值化方法
    # 自适应阈值法
    binary1 = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                   cv2.THRESH_BINARY_INV, 11, 2)
    
    # Otsu 二值化
    _, binary2 = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    
    # 尝试不同参数的二值化
    _, binary3 = cv2.threshold(gray, 180, 255, cv2.THRESH_BINARY_INV)
    
    # 使用高斯模糊去噪然后再二值化
    blurred = cv2.GaussianBlur(gray, (3, 3), 0)
    _, binary4 = cv2.threshold(blurred, 150, 255, cv2.THRESH_BINARY_INV)
    
    # 形态学操作：先腐蚀后膨胀，去除小噪点
    kernel = np.ones((2, 2), np.uint8)
    eroded = cv2.erode(binary4, kernel, iterations=1)
    dilated = cv2.dilate(eroded, kernel, iterations=1)
    
    # 显示所有处理后的图像
    if debug:
        debug_dir = "debug_images"
        os.makedirs(debug_dir, exist_ok=True)
        
        plt.figure(figsize=(15, 10))
        plt.subplot(231), plt.imshow(gray, cmap='gray'), plt.title('灰度图')
        plt.subplot(232), plt.imshow(binary1, cmap='gray'), plt.title('自适应二值化')
        plt.subplot(233), plt.imshow(binary2, cmap='gray'), plt.title('Otsu二值化')
        plt.subplot(234), plt.imshow(binary3, cmap='gray'), plt.title('固定阈值二值化')
        plt.subplot(235), plt.imshow(binary4, cmap='gray'), plt.title('高斯模糊+二值化')
        plt.subplot(236), plt.imshow(dilated, cmap='gray'), plt.title('腐蚀+膨胀')
        plt.tight_layout()
        plt.savefig(os.path.join(debug_dir, 'preprocessing_results.png'))
        print("预处理结果已保存到 debug_images/preprocessing_results.png")
    
    # 创建一个处理结果数组
    results = [binary1, binary2, binary3, binary4, dilated]
    return results


def find_best_contours(preprocessed_images):
    """
    从多种预处理结果中寻找最佳轮廓
    """
    best_contours = []
    best_count = 0
    best_image = None
    
    # 对每种处理方法找轮廓
    for i, processed in enumerate(preprocessed_images):
        contours, _ = cv2.findContours(processed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 过滤轮廓
        valid_contours = []
        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            if 5 < w < 50 and h > 10 and w < h * 3:  # 过滤出可能是数字的轮廓
                valid_contours.append(contour)
        
        print(f"预处理方法 {i+1}: 找到 {len(valid_contours)} 个有效轮廓")
        
        # 如果当前方法找到的有效轮廓最多，则更新最佳结果
        if len(valid_contours) > best_count and 3 <= len(valid_contours) <= 10:
            best_count = len(valid_contours)
            best_contours = valid_contours
            best_image = processed
    
    print(f"最佳预处理方法找到 {best_count} 个有效轮廓")
    
    # 如果没找到任何轮廓，使用备用方法
    if best_count == 0:
        print("使用备用方法: 对图像进行垂直投影分割")
        best_image = preprocessed_images[3]  # 使用高斯模糊+二值化的结果
        digit_regions = vertical_projection_segmentation(best_image)
        
        # 为每个区域创建轮廓
        best_contours = []
        for region in digit_regions:
            x, y, w, h = region
            # 创建一个矩形轮廓
            contour = np.array([[x, y], [x+w, y], [x+w, y+h], [x, y+h]]).reshape(-1, 1, 2)
            best_contours.append(contour)
        
        print(f"备用方法找到 {len(best_contours)} 个数字区域")
    
    return best_contours, best_image


def vertical_projection_segmentation(binary_image):
    """
    使用垂直投影法分割图像
    """
    # 计算垂直投影
    projection = np.sum(binary_image, axis=0)
    
    # 绘制投影图
    debug_dir = "debug_images"
    os.makedirs(debug_dir, exist_ok=True)
    
    plt.figure(figsize=(10, 5))
    plt.subplot(211), plt.imshow(binary_image, cmap='gray')
    plt.subplot(212), plt.plot(projection)
    plt.tight_layout()
    plt.savefig(os.path.join(debug_dir, 'projection.png'))
    print("垂直投影图已保存到 debug_images/projection.png")
    
    # 找出空白区域（投影值低的区域）
    threshold = np.mean(projection) * 0.3
    regions = []
    start = None
    
    for i, p in enumerate(projection):
        if start is None and p > threshold:
            start = i
        elif start is not None and p <= threshold:
            if i - start > 5:  # 忽略太窄的区域
                regions.append((start, i))
            start = None
    
    # 处理最后一个区域
    if start is not None:
        regions.append((start, len(projection)))
    
    print(f"通过垂直投影找到 {len(regions)} 个可能的数字区域")
    
    # 将区域转换为 (x, y, w, h) 格式
    digit_regions = []
    for start, end in regions:
        digit_regions.append((start, 0, end - start, binary_image.shape[0]))
    
    return digit_regions


def extract_digits_from_contours(contours, original_image):
    """
    从轮廓中提取数字
    """
    debug_dir = "debug_images"
    os.makedirs(debug_dir, exist_ok=True)
    
    # 创建一个彩色图像用于显示轮廓
    contour_image = cv2.cvtColor(original_image.copy(), cv2.COLOR_GRAY2BGR)
    
    digits = []
    # 根据轮廓位置排序（从左到右）
    contours = sorted(contours, key=lambda c: cv2.boundingRect(c)[0])
    
    for i, contour in enumerate(contours):
        x, y, w, h = cv2.boundingRect(contour)
        
        # 在彩色图像上绘制轮廓和序号
        cv2.drawContours(contour_image, [contour], 0, (0, 255, 0), 2)
        cv2.putText(contour_image, str(i), (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 2)
        
        # 提取可能的数字图像
        digit_img = original_image[y:y+h, x:x+w]
        
        # 计算特征
        aspect_ratio = w / float(h) if h > 0 else 0
        area = cv2.contourArea(contour)
        
        # 存储数字及其特征
        digit_info = {
            'img': digit_img,
            'x': x,
            'y': y,
            'w': w,
            'h': h,
            'area': area,
            'aspect_ratio': aspect_ratio,
            'pixel_sum': np.sum(digit_img) / 255.0,
        }
        digits.append(digit_info)
    
    # 保存带有轮廓的图像
    cv2.imwrite(os.path.join(debug_dir, 'contours.png'), contour_image)
    print("轮廓标记图像已保存到 debug_images/contours.png")
    
    return digits


def recognize_digits(digits):
    """
    识别提取的数字
    """
    result = ""
    
    debug_dir = "debug_images"
    os.makedirs(debug_dir, exist_ok=True)
    
    # 保存每个数字的图像用于调试
    plt.figure(figsize=(15, 3))
    
    # 简单的基于规则的识别
    for i, digit in enumerate(digits):
        img = digit['img']
        h, w = img.shape
        aspect_ratio = digit['aspect_ratio']
        pixel_sum = digit['pixel_sum']
        density = pixel_sum / (h * w) if h * w > 0 else 0
        
        # 输出调试信息
        print(f"数字 {i}: 高={h}, 宽={w}, 宽高比={aspect_ratio:.2f}, 像素密度={density:.2f}")
        
        # 显示数字图像
        if i < 10:  # 最多显示10个数字
            plt.subplot(1, 10, i+1)
            plt.imshow(img, cmap='gray')
            plt.axis('off')
        
        # 基于特征的识别逻辑
        recognized = '?'  # 默认为未识别
        
        # 以下逻辑需要根据实际验证码类型调整
        if aspect_ratio < 0.5:  # 细长的可能是1
            recognized = "1"
        elif 0.5 <= aspect_ratio < 0.7:
            if density < 0.5:
                recognized = "7"
            else:
                recognized = "3"
        elif 0.7 <= aspect_ratio < 0.9:
            if density < 0.5:
                recognized = "4"
            elif density < 0.7:
                recognized = "2"
            else:
                recognized = "8"
        else:  # 宽高比较大
            if density < 0.5:
                recognized = "9"
            elif density < 0.7:
                recognized = "6"
            else:
                recognized = "0"
        
        result += recognized
        print(f"识别为: {recognized}")
    
    # 保存数字图像
    plt.tight_layout()
    plt.savefig(os.path.join(debug_dir, 'digits.png'))
    print("分割的数字图像已保存到 debug_images/digits.png")
    
    return result


def auto_recognize_captcha(image_path, allow_manual_correction=True):
    """
    自动识别验证码图片中的数字
    
    Args:
        image_path: 验证码图片路径
        allow_manual_correction: 是否允许手动纠正识别结果
        
    Returns:
        识别出的验证码文本
    """
    try:
        # 创建调试图像目录
        debug_dir = "debug_images"
        os.makedirs(debug_dir, exist_ok=True)
        
        # 保存原始图像到调试目录
        orig_img = cv2.imread(image_path)
        cv2.imwrite(os.path.join(debug_dir, 'original.png'), orig_img)
        
        # 预处理图像获得多个结果
        preprocessed_images = preprocess_captcha_image(image_path)
        
        # 找到最佳轮廓
        best_contours, best_image = find_best_contours(preprocessed_images)
        
        # 提取数字
        digits = extract_digits_from_contours(best_contours, best_image)
        
        # 识别数字
        result = recognize_digits(digits)
        
        print(f"自动识别结果: {result}")
        
        # 如果允许手动纠正，且识别出了内容
        if allow_manual_correction and len(digits) > 0:
            # 显示原始验证码
            print("\n请查看 debug_images 目录下的图片，特别是 original.png 和 digits.png")
            print("自动识别可能不准确，请输入正确的验证码（如果自动识别正确，直接按回车）:")
            
            manual_input = input()
            if manual_input:
                result = manual_input
                print(f"使用手动输入的验证码: {result}")
        
        return result
    except Exception as e:
        print(f"自动识别验证码失败: {e}")
        import traceback
        traceback.print_exc()
        
        if allow_manual_correction:
            print("\n自动识别失败，请查看原始验证码图片并手动输入:")
            manual_result = input()
            return manual_result
        
        return ""


def save_captcha_model(image_path, correct_value):
    """
    保存验证码图片和对应的正确值，可用于后续训练模型
    """
    try:
        # 创建样本目录
        samples_dir = "captcha_samples"
        os.makedirs(samples_dir, exist_ok=True)
        
        # 复制图片到样本目录，文件名为正确的值
        img = cv2.imread(image_path)
        sample_path = os.path.join(samples_dir, f"{correct_value}.png")
        cv2.imwrite(sample_path, img)
        
        print(f"已保存验证码样本: {sample_path}")
        return True
    except Exception as e:
        print(f"保存验证码样本失败: {e}")
        return False


if __name__ == '__main__':
    # 测试验证码识别
    import sys
    if len(sys.argv) > 1:
        image_path = sys.argv[1]
    else:
        image_path = "captcha.png"
    
    result = auto_recognize_captcha(image_path)
    print(f"识别结果: {result}")
    
    # 如果有识别结果，询问是否正确并保存样本
    if result:
        print("\n识别结果是否正确? (y/n)")
        is_correct = input().lower() == 'y'
        
        if is_correct:
            save_captcha_model(image_path, result)
        else:
            print("请输入正确的验证码值:")
            correct_value = input()
            if correct_value:
                save_captcha_model(image_path, correct_value) 